import os
import subprocess
import argparse
import re

def detect_gpu_type():
    """检测系统可用的GPU类型"""
    try:
        # 检测NVIDIA GPU
        subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return 'nvidia'
    except (subprocess.CalledProcessError, FileNotFoundError):
        try:
            # 检测AMD GPU
            subprocess.run(['rocminfo'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            return 'amd'
        except (subprocess.CalledProcessError, FileNotFoundError):
            try:
                # 检测Intel GPU
                subprocess.run(['intel_gpu_top'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
                return 'intel'
            except (subprocess.CalledProcessError, FileNotFoundError):
                return None

def get_gpu_encoder(gpu_type):
    """根据GPU类型返回最佳H.265编码器"""
    if gpu_type == 'nvidia':
        return 'hevc_nvenc'  # NVIDIA H.265编码器
    elif gpu_type == 'amd':
        return 'hevc_amf'    # AMD H.265编码器
    elif gpu_type == 'intel':
        return 'hevc_qsv'    # Intel H.265编码器
    else:
        return 'libx265'     # CPU H.265编码器

def convert_mkv_to_mp4(input_file, output_file=None, overwrite=False, gpu_accel=True, quality='high'):
    """
    将MKV文件转换为MP4格式，使用GPU加速并优化质量
    
    参数:
    input_file: 输入的MKV文件路径
    output_file: 输出的MP4文件路径（默认为输入文件同目录同名）
    overwrite: 是否覆盖已存在的输出文件
    gpu_accel: 是否使用GPU加速
    quality: 视频质量设置（'low', 'medium', 'high'）
    """
    # 设置默认输出文件名
    if output_file is None:
        filename, _ = os.path.splitext(input_file)
        output_file = filename + ".mp4"
    
    # 检查输入文件是否存在
    if not os.path.isfile(input_file):
        raise FileNotFoundError(f"输入文件不存在: {input_file}")
    
    # 删除上一次的输出文件
    if os.path.isfile(output_file):
        if overwrite:
            os.remove(output_file)
        else:
            print(f"错误: 输出文件已存在 - {output_file} (使用 -y 选项覆盖)")
            return False
    
    # 检测GPU类型
    gpu_type = detect_gpu_type() if gpu_accel else None
    encoder = get_gpu_encoder(gpu_type) if gpu_accel else 'libx265'
    
    print(f"使用编码器: {encoder} ({gpu_type or 'CPU'})")
    print(f"质量设置: {quality}")
    
    # 构建FFmpeg命令 - 优化质量版
    cmd = [
        'ffmpeg',
        '-i', input_file,
        '-c:v', encoder,  # 使用GPU加速编码器
        '-c:a', 'copy',    # 拷贝音频流
        '-movflags', '+faststart',  # 优化网络播放
        '-y' if overwrite else '-n',  # 自动覆盖/不覆盖
    ]
    
    # 根据质量设置添加参数 - 优化版
    if quality == 'high':
        # 高质量设置（优化码率）
        if encoder == 'hevc_nvenc':
            cmd.extend(['-preset', 'p7', '-tune', 'hq', '-cq', '23', '-b:v', '0'])
        elif encoder == 'hevc_amf':
            cmd.extend(['-quality', 'quality', '-qp_i', '23', '-qp_p', '23', '-b:v', '0'])
        elif encoder == 'hevc_qsv':
            cmd.extend(['-preset', 'slower', '-global_quality', '23', '-b:v', '0'])
        else:  # CPU编码
            # 修复: 使用x265参数格式
            cmd.extend(['-preset', 'slow', '-x265-params', 'crf=23:psy-rd=1.0:aq-mode=2'])
    elif quality == 'medium':
        # 中等质量（平衡）
        if encoder == 'hevc_nvenc':
            cmd.extend(['-preset', 'p5', '-cq', '26', '-b:v', '0'])
        elif encoder == 'hevc_amf':
            cmd.extend(['-quality', 'balanced', '-qp_i', '26', '-qp_p', '26', '-b:v', '0'])
        elif encoder == 'hevc_qsv':
            cmd.extend(['-preset', 'medium', '-global_quality', '26', '-b:v', '0'])
        else:  # CPU编码
            # 修复: 使用x265参数格式
            cmd.extend(['-preset', 'medium', '-x265-params', 'crf=26:psy-rd=0.8:aq-mode=1'])
    else:  # low
        # 较低质量（文件较小）
        if encoder == 'hevc_nvenc':
            cmd.extend(['-preset', 'p3', '-cq', '28', '-b:v', '0'])
        elif encoder == 'hevc_amf':
            cmd.extend(['-quality', 'speed', '-qp_i', '28', '-qp_p', '28', '-b:v', '0'])
        elif encoder == 'hevc_qsv':
            cmd.extend(['-preset', 'fast', '-global_quality', '28', '-b:v', '0'])
        else:  # CPU编码
            # 修复: 使用x265参数格式
            cmd.extend(['-preset', 'fast', '-x265-params', 'crf=28:psy-rd=0.5:aq-mode=1'])
    
    cmd.append(output_file)
    
    # 打印完整命令用于调试
    print("完整FFmpeg命令:")
    print(" ".join(cmd))
    
    try:
        # 执行转换命令
        print("开始转换...")
        process = subprocess.Popen(
            cmd,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            encoding='utf-8',
            errors='replace'
        )
        
        # 实时显示进度
        while True:
            line = process.stderr.readline()
            if not line:
                break
            # 提取并显示进度信息
            if 'time=' in line:
                time_match = re.search(r'time=(\d+:\d+:\d+\.\d+)', line)
                if time_match:
                    print(f"\r转换进度: {time_match.group(1)}", end='', flush=True)
            # 打印错误信息
            if 'error' in line.lower() or 'failed' in line.lower():
                print(f"\n错误信息: {line.strip()}")
        
        # 等待进程完成
        process.wait()
        
        if process.returncode == 0:
            print(f"\n转换成功: {input_file} -> {output_file}")
            return True
        else:
            print(f"\n转换失败，返回码: {process.returncode}")
            # 打印错误信息
            print("错误信息:")
            # 读取剩余的错误输出
            while True:
                line = process.stderr.readline()
                if not line:
                    break
                print(line, end='')
            return False
    except FileNotFoundError:
        print("\n错误: 未找到FFmpeg，请先安装FFmpeg并添加到系统路径")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='使用GPU加速将MKV文件转换为MP4格式并优化质量（H.265编码）')
    parser.add_argument('input', help='输入MKV文件路径')
    parser.add_argument('-o', '--output', help='输出MP4文件路径（可选）')
    parser.add_argument('-y', '--overwrite', action='store_true', help='覆盖已存在的输出文件')
    parser.add_argument('--no-gpu', action='store_true', help='禁用GPU加速（使用CPU编码）')
    parser.add_argument('-q', '--quality', choices=['low', 'medium', 'high'], default='high',
                        help='视频质量设置（low, medium, high），默认为high')
    
    args = parser.parse_args()
    
    convert_mkv_to_mp4(
        input_file=args.input,
        output_file=args.output,
        overwrite=args.overwrite,
        gpu_accel=not args.no_gpu,
        quality=args.quality
    )